# get_cost.py
# 仅把 Excel 操作改为 CSV，其余逻辑不变
import csv
import math
import os
import time
from datetime import datetime
from pathlib import Path
import numpy as np
from simpy import Environment

from osc05.resource.MonitoredFilterStore import MonitoredFilterStore
from osc05.resource.Resources import ResourceUtil
from osc05.resource.stakeholders import ConstructionSite
from osc05.static.constants import Constants, Debugs, DirectoryConfig
from osc05.utils.cpm_class import ProjectScheduler
from osc05.utils.random_util import RandomUtil
from osc05.utils.calibration_random_util import CalibrationRandomUtil  # 新增导入


# -----------------------------------------------------------
# 1. 主入口：get_cost_with_simulation
# -----------------------------------------------------------
def get_cost_with_simulation(
        start_correction,
        seed=None,
        unit_waiting_penalty=Debugs.unit_component_late_acceptance_penalty_by_day,
        unit_late_delivery_penalty=Debugs.unit_late_delivery_penalty_by_day,
        unit_project_late_completion_penalty=Debugs.unit_project_late_completion_penalty_by_day,
        project_deadline=Debugs.project_deadline,
        sample_size=None,  # 新增：支持自定义样本量
        calibration_seed=None,  # 新增：校准种子参数
        debug=False):
    now = datetime.now()

    # 确定实际使用的样本量
    if sample_size is None:
        actual_sample_size = Constants.SCENARIO_NUM
    else:
        actual_sample_size = sample_size

    scheduler = ProjectScheduler()
    activities, prj_start_date = scheduler.get_schedule('2025-3-10')

    # 确定使用的种子：优先使用 calibration_seed，其次使用 seed
    actual_seed = calibration_seed if calibration_seed is not None else seed

    samples_dir = Path(__file__).parent.parent / "sampling" / "samples"

    # 关键修改：如果提供了校准参数，使用支持校准的 RandomUtil
    if sample_size is not None or calibration_seed is not None:
        # 使用支持校准的 RandomUtil
        random_util = CalibrationRandomUtil(
            seed=actual_seed,
            n_sf=Constants.N_ACTIVITY,
            samples_dir=samples_dir,
            sample_size=sample_size,
            calibration_seed=calibration_seed
        )
        print(f"使用校准模式: 样本量={sample_size}, 校准种子={calibration_seed}")
    else:
        # 使用原有的 RandomUtil（保持向后兼容）
        random_util = RandomUtil(actual_seed, n_sf=Constants.N_ACTIVITY, samples_dir=samples_dir)

    sum_additional_cost_every_scenario = 0
    total_additional_cost_every_scenario_in_detail = [0, 0, 0]

    if debug:
        # 1) 明细文件
        detail_path = os.path.join(DirectoryConfig.OUTPUT_DIR, 'x_and_cost.csv')
        file_exists = os.path.isfile(detail_path)
        with open(detail_path, 'a', newline='', encoding='utf-8') as f_detail:
            writer_detail = csv.writer(f_detail)
            if not file_exists:
                writer_detail.writerow(
                    ["Date", "Hr", "Min", "样本量", "种子", "场景编号", "total_waiting_cost",
                     "total_late_delivery_penalty", "project_late_completion_penalty", "total_cost"] +
                    [f'FL{i:02d}' for i in range(1, Constants.N_ACTIVITY + 1)]
                )

            cost_list = []
            for i in range(actual_sample_size):
                total_waiting_cost, total_late_delivery_penalty, project_late_completion_penalty = \
                    simulate(activities, debug, project_deadline, i + 1, start_correction,
                             unit_late_delivery_penalty,
                             unit_project_late_completion_penalty, unit_waiting_penalty,
                             random_util=random_util)
                total = total_waiting_cost - total_late_delivery_penalty + project_late_completion_penalty

                # 记录校准参数信息
                sample_size_info = actual_sample_size if sample_size is not None else "default"
                seed_info = actual_seed if actual_seed is not None else "default"

                print(f"样本量={sample_size_info}, 种子={seed_info}, 场景={i}, 总成本={total}")

                writer_detail.writerow([
                                           now.strftime("%y-%m-%d"), now.strftime("%H"), now.strftime("%M"),
                                           sample_size_info, seed_info, i + 1,
                                           total_waiting_cost, -total_late_delivery_penalty,
                                           project_late_completion_penalty, total
                                       ] + list(start_correction))

                cost_list.append([total_waiting_cost, total_late_delivery_penalty,
                                  project_late_completion_penalty, total])

                total_additional_cost_every_scenario_in_detail = [
                    m + n for m, n in zip(
                        total_additional_cost_every_scenario_in_detail,
                        [total_waiting_cost, total_late_delivery_penalty, project_late_completion_penalty])
                ]

            # 2) 汇总文件
            summary_path = os.path.join(DirectoryConfig.OUTPUT_DIR, 'x_and_cost_summary.csv')
            file_exists_sum = os.path.isfile(summary_path)
            with open(summary_path, 'a', newline='', encoding='utf-8') as f_sum:
                writer_sum = csv.writer(f_sum)
                if not file_exists_sum:
                    writer_sum.writerow([
                                            "Date", "Hr", "Min", "样本量", "种子",
                                            "max_total", "min_total", "mean_total",
                                            "std_total", "Coeff_of_Var", "max_project_delay_days"
                                        ] + [f'FL{i:02d}' for i in range(1, Constants.N_ACTIVITY + 1)])

                np_array = np.array(cost_list)
                max_total = np.max(np_array[:, 3])
                min_total = np.min(np_array[:, 3])
                mean_total = np.mean(np_array[:, 3])
                std_total = np.std(np_array[:, 3])
                max_project_delay_days = np.max(np_array[:, 2]) / unit_project_late_completion_penalty

                summary_row = [
                                  now.strftime("%y-%m-%d"), now.strftime("%H"), now.strftime("%M"),
                                  sample_size_info, seed_info,
                                  max_total, min_total, mean_total, std_total,
                                  f"{std_total / mean_total:.1f}", max_project_delay_days
                              ] + list(start_correction)
                writer_sum.writerow(summary_row)

                # 写平均行
                avg_row = [
                              now.strftime("%y-%m-%d"), now.strftime("%H"), now.strftime("%M"),
                              "AVE", sample_size_info, seed_info,
                              total_additional_cost_every_scenario_in_detail[0] / actual_sample_size,
                              -total_additional_cost_every_scenario_in_detail[1] / actual_sample_size,
                              total_additional_cost_every_scenario_in_detail[2] / actual_sample_size,
                              (total_additional_cost_every_scenario_in_detail[0] -
                               total_additional_cost_every_scenario_in_detail[1] +
                               total_additional_cost_every_scenario_in_detail[2]) / actual_sample_size
                          ] + list(start_correction)
                writer_detail.writerow(avg_row)
                writer_detail.writerow(start_correction)

                print("Debug=True, average total=",
                      (total_additional_cost_every_scenario_in_detail[0] -
                       total_additional_cost_every_scenario_in_detail[1] +
                       total_additional_cost_every_scenario_in_detail[2]) / actual_sample_size)

        return (
            total_additional_cost_every_scenario_in_detail[0] / actual_sample_size,
            total_additional_cost_every_scenario_in_detail[1] / actual_sample_size,
            total_additional_cost_every_scenario_in_detail[2] / actual_sample_size
        )

    else:
        # 非 debug 分支
        for i in range(actual_sample_size):
            additional_cost = simulate(
                activities, debug, project_deadline, i + 1, start_correction,
                unit_late_delivery_penalty,
                unit_project_late_completion_penalty, unit_waiting_penalty,
                random_util)
            sum_additional_cost_every_scenario += additional_cost
        return sum_additional_cost_every_scenario / actual_sample_size


# -----------------------------------------------------------
# 2. 模拟函数：simulate
# -----------------------------------------------------------
def simulate(activities, debug, project_deadline, seed, start_correction,
             unit_late_delivery_penalty,
             unit_project_late_completion_penalty, unit_waiting_penalty, random_util):
    env = Environment()
    storage_yard = MonitoredFilterStore(env, capacity=200)
    assembly_worker_teams = ResourceUtil.get_stations(env, capacity=1)
    grout_gears = ResourceUtil.get_stations(env, capacity=1)
    construction_site1 = ConstructionSite(
        env=env, name='Site1',
        storage_yard=storage_yard,
        worker_teams=assembly_worker_teams,
        grout_gears=grout_gears,
        crane_num=1,
        activities=activities)
    # 可以在外面创建对象，循环开始初始化
    smp_delivery_delay, smp_duration = get_sampled_duration_and_delivery_time(random_util, seed)

    # 设置活动参数
    for idx, activity in enumerate(activities):
        activity.construction_site = construction_site1
        activity.completion_event = env.event()
        activity.component_delivered_event = env.event()
        activity.act_duration = math.ceil(smp_duration[idx])
        activity.skd_delivery_time = max(activity.es + start_correction[idx], 0)
        activity.act_delivery_time = activity.skd_delivery_time + max(0, math.ceil(smp_delivery_delay[idx]))

    # 启动仿真进程
    for activity in activities[:96]:
        env.process(activity.transport())
    for activity in activities[:96]:
        env.process(activity.act())
    env.run(until=2000)

    # 计算成本
    total_waiting_cost = 0
    total_late_delivery_penalty = 0
    project_act_completion = 0

    if debug:
        csv_name = f'Project_process_simulation_{seed}.csv'
        csv_path = os.path.join(DirectoryConfig.OUTPUT_DIR, csv_name)
        with open(csv_path, 'w', newline='', encoding='utf-8') as f:
            writer = csv.writer(f)
            writer.writerow([
                'name', 'es', 'duration', 'chromo', 'skd_delivery_time', 'act_delivery_time',
                'act_erection_start', 'act_duration', 'delivery_delay',
                'start_delay', 'waiting_cost', 'penalty_cmp_delivery', 'penalty_prj_completion'
            ])

    for idx, activity in enumerate(activities):
        total_waiting_cost += unit_waiting_penalty * (activity.act_erection_start - activity.act_delivery_time)
        total_late_delivery_penalty += unit_late_delivery_penalty * max(
            0, (activity.act_delivery_time - activity.skd_delivery_time))
        current_act_completion = activity.act_erection_start + activity.act_duration
        project_act_completion = max(project_act_completion, current_act_completion)

        if debug:
            row = [
                activity.name, round(activity.es, 2), activity.duration,
                round(start_correction[idx]),
                round(activity.skd_delivery_time, 2), round(activity.act_delivery_time, 2),
                round(activity.act_erection_start, 2), round(activity.act_duration, 2),
                round(activity.act_delivery_time - activity.skd_delivery_time, 2),
                round(activity.act_erection_start - activity.es, 2),
                round(unit_waiting_penalty * (activity.act_erection_start - activity.act_delivery_time), 2),
                round(unit_late_delivery_penalty * max(0, activity.act_delivery_time - activity.skd_delivery_time), 2)
            ]
            with open(csv_path, 'a', newline='', encoding='utf-8') as f:
                csv.writer(f).writerow(row)

    project_late_completion_penalty = max(0,
                                          project_act_completion - project_deadline) * unit_project_late_completion_penalty
    total_cost = total_waiting_cost - total_late_delivery_penalty + project_late_completion_penalty

    if debug:
        summary_row = [''] * 8 + [
            'T.Cost', total_cost,
            'Sum_waiting', total_waiting_cost,
            'Sum_late_delivery', total_late_delivery_penalty,
            'Prj_delay_days', max(0, project_act_completion - project_deadline)
        ]
        with open(csv_path, 'a', newline='', encoding='utf-8') as f:
            csv.writer(f).writerow(summary_row)

    if debug:
        return total_waiting_cost, total_late_delivery_penalty, project_late_completion_penalty
    else:
        return total_cost


# -----------------------------------------------------------
# 3. 其它辅助函数
# -----------------------------------------------------------
def get_sampled_duration_and_delivery_time(random_util, seed=None) -> tuple[list[float], list[float]]:
    (smp_durations,
     smp_delivery_delay_days,
     smp_crane_failure_times,
     smp_rain_hold_daily_events,
     smp_strong_wind_hold_daily_events) = random_util.get_samples()
    return smp_delivery_delay_days, smp_durations


# -----------------------------------------------------------
# 4. 本地测试入口（保持不变）
# -----------------------------------------------------------
if __name__ == '__main__':
    print('base plan')
    delivery_time_offset = [0] * Constants.N_ACTIVITY
    for seed in range(1, 31):
        total_waiting_cost = get_cost_with_simulation(
            start_correction=delivery_time_offset,
            seed=seed,
            debug=True)
        print(total_waiting_cost)
    print('####')